Skip to content

Optimize truncateFileTree token estimation and sampling seed - #1216

Open
nordicnode wants to merge 2 commits into
CodebuffAI:mainfrom
nordicnode:perf-truncate-file-tree-token-budget
Open

Optimize truncateFileTree token estimation and sampling seed#1216
nordicnode wants to merge 2 commits into
CodebuffAI:mainfrom
nordicnode:perf-truncate-file-tree-token-budget

Conversation

@nordicnode

@nordicnode nordicnode commented Sep 2, 2026

Copy link
Copy Markdown

Summary

  • In packages/agent-runtime/src/system-prompt/truncate-file-tree.ts, optimize token counting, sampling seed generation, and unimportant file filtering during prompt file tree truncation.
  • Deterministic lightweight seed: Previously, sampleSizeWithSeed was passed JSON.stringify(sortedFiles) + JSON.stringify(sampleCount) as its seed. On repositories with thousands of files, serializing the entire array of tree node objects created a multi-megabyte string, costing ~68ms per 50 runs. Replaced this with a lightweight deterministic seed (${sortedFiles.length}:${sampleCount}:${sortedFiles[0]?.path ?? ''}:${sortedFiles[sortedFiles.length - 1]?.path ?? ''}) that computes in 0.01ms (over 6,800x faster).
  • Accurate raw-string token counting: Replaced countTokensJson(printedTree) with countTokens(printedTree). printedTree is directly interpolated into markdown template literals in the system prompt (system-prompt/prompts.ts), never JSON-encoded. Calling countTokensJson executed JSON.stringify(str) to escape quotes and newlines, unnecessarily inflating token counts and allocating intermediate strings. countTokens directly measures the true prompt token cost and runs 3.2x faster (6.71ms vs 21.31ms).
  • Pre-partitioned unimportant filters: Partitioned unimportantExtensions into UNIMPORTANT_DIR_PATTERNS and UNIMPORTANT_EXTENSIONS so directory and file filters do not repeatedly evaluate ext.startsWith('/') across 51 items for every single file in the tree.
  • Preserved .so binary filtering: Confirmed and restored .so alongside .exe, .dll, and .lib in UNIMPORTANT_EXTENSIONS.
  • Added unit tests: Added comprehensive unit tests in packages/agent-runtime/src/system-prompt/__tests__/truncate-file-tree.test.ts covering budget thresholds, build directory filtering, depth-based fallback, and compiled binary/shared library (.so, .dll, .exe, .lib) filtering.

Test plan

  • bun test src/system-prompt/__tests__/truncate-file-tree.test.ts (4 passed, 0 failed)
  • bun test src/tools/handlers/__tests__/read-subtree.test.ts (7 passed, 0 failed)
  • bun run build:sdk (successful build)
  • bun freebuff/cli/build.ts 0.0.0-ci (successful binary build)
  • Binary smoke test: bun cli/scripts/smoke-binary.ts cli/bin/freebuff (OK)
  • Prettier check: bun x prettier --check packages/agent-runtime/src/system-prompt/truncate-file-tree.ts packages/agent-runtime/src/system-prompt/__tests__/truncate-file-tree.test.ts (All matched files use Prettier code style)

@codebuff-team

Copy link
Copy Markdown
Contributor

Nice diagnosis on the seed and countTokensJson issues — using JSON.stringify(sortedFiles) purely as a random seed was wasteful, and if printedTree is genuinely consumed as a raw string downstream (not JSON-encoded), switching to countTokens is a legitimate accuracy + perf win. The dir/extension list split (UNIMPORTANT_DIR_PATTERNS / UNIMPORTANT_EXTENSIONS) is a clean way to avoid the startsWith('/') check per file per extension.

However, in the extension list refactor at the bottom of truncate-file-tree.ts, .so disappears entirely — it's present in the old unimportantExtensions array (grouped with .exe, .dll, .lib) but missing from the new UNIMPORTANT_EXTENSIONS. That's an unannounced behavior change: shared libraries will now show up in the file tree where they were previously filtered. If intentional, call it out in the PR description; if accidental (looks like it), please restore it and add a regression test — your new test suite covers .min.js filtering but nothing for the executable/library category, so this slipped through untested.

Separately, the countTokensJsoncountTokens swap changes token-budget behavior across the whole truncation pipeline (four call sites). That's a bigger functional change than the PR title suggests ('optimize... estimation') — it's really a correctness change to what gets counted. Worth confirming with whoever owns token-counter.ts that printedTree was never intended to model JSON-escaped cost, since if it was, this now under-counts tokens and could cause budget overruns in production.

Good test additions overall; would like one covering the extension-list edge case before this goes in.

@codebuff-team codebuff-team added bot:triaged Classified by the community triage bot pr:needs-work Right idea, not mergeable as written labels Sep 3, 2026
@nordicnode

Copy link
Copy Markdown
Author

Thanks for the keen eye @codebuff-team! Both points have been addressed in the latest commit:

  1. Restored .so Extension & Added Regression Test:

    • The omission of .so from UNIMPORTANT_EXTENSIONS was indeed accidental during the list split. It has been restored alongside .exe, .dll, and .lib.
    • Added a regression test in truncate-file-tree.test.ts verifying that compiled binary and library files (.so, .dll, .exe, .lib) are filtered out by removeUnimportantFiles.
  2. countTokens vs countTokensJson Rationale:

    • Verified the downstream usage in packages/agent-runtime/src/system-prompt/prompts.ts: printedTree is directly interpolated into raw markdown template literals:
      <file_tree>
      ${printedTree}
      </file_tree>
      It is never serialized as JSON.
    • Calling countTokensJson was running JSON.stringify(string) under the hood, which wrapped the string in double quotes and escaped every newline (\n -> \\n), artificially inflating token estimates and spending time copying strings. countTokens(printedTree) reflects the true token count of the raw string sent to the model.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bot:triaged Classified by the community triage bot pr:needs-work Right idea, not mergeable as written

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants